"""
基于论文[1] Christian J. A. , Derksen H. , Watkins R. .Lunar crater identification in digital Images[J/OL].J. Astronaut. Sci.,2021,68(4):1056-1144

仅使用论文提供的不变量作为描述子，由于描述子没有顺序，因此只能使用投票法。
"""

import numpy as np
import torch
from utils.ellipse import radius_ellipse
from .base import TriadPyramidVote


class ChristianPyramidVote(TriadPyramidVote):
    def __init__(self, catalog_path, device="cuda:0", limit_N=5, **kwargs):
        super().__init__(catalog_path, device, **kwargs)
        self.limit_N = limit_N

    def triad_descriptor(
        self, i, j, k, Q1, Q2, Q3, *args, th=0.05, factor=0.3, **kwargs
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """从陨石坑的参数中计算描述子，C1和C2可以是含有不确定度的椭圆参数"""
        N = Q3.shape[0]
        if len(Q1.shape) == 2:
            Q1 = Q1[None].repeat(N, dim=0)
        if len(Q2.shape) == 2:
            Q2 = Q2[None].repeat(N, dim=0)
        assert Q1.shape[0] == Q2.shape[0] == Q3.shape[0]
        d1 = np.mean(radius_ellipse(Q1), axis=0)
        d2 = np.mean(radius_ellipse(Q2), axis=0)
        d3 = np.mean(radius_ellipse(Q3), axis=0)
        # 按直径从小到大排序
        D = np.array((d1, d2, d3))
        # 任取两个直径计算比值，当比值接近1的数超过2个时，排除
        valid_ind = np.abs(D[[1, 2]] / D[0, None] - 1) > th
        valid_ind &= np.abs(D[[0, 2]] / D[1, None] - 1) > th
        valid_ind &= np.abs(D[[0, 1]] / D[2, None] - 1) > th
        valid_ind = valid_ind.all(axis=0)
        index = np.argsort((d1, d2, d3), axis=0)
        Q = np.array((Q1, Q2, Q3))[index, np.arange(N)]
        Q = torch.tensor(Q, device=self.device, dtype=torch.float32)
        # 计算不变量
        I1 = torch.einsum("nii->n", Q[0] @ torch.inverse(Q[1])) * torch.einsum(
            "nii->n", Q[1] @ torch.inverse(Q[0])
        )
        I2 = torch.einsum("nii->n", Q[1] @ torch.inverse(Q[2])) * torch.einsum(
            "nii->n", Q[2] @ torch.inverse(Q[1])
        )
        I3 = torch.einsum("nii->n", Q[2] @ torch.inverse(Q[0])) * torch.einsum(
            "nii->n", Q[0] @ torch.inverse(Q[2])
        )
        # 每行是一个3x6的点列，每三个元素代表一个点
        I = torch.stack((I1, I2, I3))
        # 按直径从小到大排序
        th = factor * factor * I.abs()
        ijk = np.vstack((i, j, k))
        return ijk, I, th, np.argsort(index, axis=0), valid_ind

    def vote_record(self, ijk, I123, th, index, vote_th=6, N=20):
        return super().vote_record(ijk, I123, th, index, vote_th, N=self.limit_N)
